rm(list = ls())
library(tidyverse)
library(lubridate)
# NOTE: INSTALL THE ABOVE PACKAGES USING THE FOLLOWING COMMANDS IF NOT INSTALLED
# install.packages("tidyverse")
# install.packages("lubridate")

if(Sys.info()["user"]=="XXXX") setwd("XXXXX")    # UPDATE THIS LINE TO POINT TO THE "Replication_Output" Directory
if(Sys.info()["user"]=="jimmy") setwd("~/Dropbox/SMS_ADMIN/Replication Package v1.0.8/Replication_Output/")


#D vector based on shares
ProportionalAssignment = function(Shares, Nt) {
  k = length(Shares)
  n_floor = floor(Nt * Shares) #number of units assigned to each treatment, rounded down
  remainder = Nt * Shares - n_floor
  units_remaining = Nt - sum(n_floor) #remaining number of units to be assigned
  
  if (units_remaining > 0) {
    Dt = c(rep(1:k, n_floor),
           # assigning each treatment acoording to rounded down average count
           sample(
             1:k,
             size = units_remaining,
             replace = T,
             prob = remainder
           )) #remaining units assigned randomly from remain replicate assignments
  } else {
    Dt = rep(1:k, n_floor)
  }
  sample(Dt)
}


DtchoiceThompson = function(Y, D, #outcomes and treatments thus far
                            k, #number of treatments
                            C, #vector of treatment cost
                            Nt) {
  # number of observations for period t
  
  A = 1 + tapply(Y, D, sum, default = 0) 
  B = 1 + tapply(1-Y, D, sum, default = 0) 
  
  Dt = rep(0, Nt)
  previousD = -Inf # auxiliary variable to avoid repeat assignments of same D
  
  for (i  in 1:Nt) {
    thetadraw = sapply(1:k, function(j)
      rbeta(1, A[j], B[j]))
    Dt[i] = which.max(thetadraw - C)
    previousD = Dt[i]
  }
  
  factor(Dt, levels = 1:k)
}



DtchoiceThompsonProbabilities = function(Y, D, #outcomes and treatments thus far
                                         k, #number of treatments
                                         C = rep(0, k), #vector of treatment cost
                                         RR = 50000) {
  #number of replication draws
  
  # Repeat Thompson sampling RR times
  DtRR = DtchoiceThompson(Y, D, k, C, RR)
  P_Dt = table(DtRR) / RR #average count for each treatment value and covariate value, replicated sample
  
  P_Dt = as_tibble(matrix(P_Dt, 1, k))
  colnames(P_Dt) = paste(1:k)
  P_Dt
}


DtchoiceThompson_modified = function(alpha) {
  if (max(alpha) < 1) {
    Shares = alpha * (1 - alpha) # Based on stationary distribution for modified Thompson
    Shares = Shares / sum(Shares)
  } else {
    Shares = alpha
  }
  
  return(Shares)
}

#priorname = paste("OUTPUT/COVID/Priors/", Sys.Date(), "priors.csv", sep = "")
#priorname = paste("OUTPUT/COVID/Priors/2020-10-15priors.csv", sep = "")
priordata = read.csv(file = "Randomization_Inference/reassignedpriors_round.csv")  %>% 
  #filter(!is.na(outcome) & !is.na(treatment) & treatment != 0) %>%
  mutate(treatment=factor(treatment))

by_treatment= priordata %>% 
  group_by(treatment, .drop=FALSE) %>% 
  summarise(avg=mean(outcome), count=n(), successes=as.integer(sum(outcome)))
alpha = DtchoiceThompsonProbabilities(
  priordata$outcome,
  priordata$treatment,
  k=length(levels(priordata$treatment)),
) %>% as_vector

P_current = DtchoiceThompson_modified(alpha)
P_current_tibble =  tibble(share = as_vector(P_current), treatment = factor(levels(priordata$treatment)))
treatment_assignment = tibble(treatment=ProportionalAssignment(P_current, 100000))
#filename = paste("OUTPUT/COVID/Treatment_assignments/", Sys.Date(), "_t_assignments.csv", sep = "")
#write_csv(treatment_assignment, filename)
write.csv(treatment_assignment,"Randomization_Inference/adaptivetreatment_round.csv")

Bayes_table=by_treatment %>%
  mutate(alpha_d= 1 + successes,
         beta= 1 + count-successes,
         mean = alpha_d / (alpha_d+beta),
         var = alpha_d * beta / ((alpha_d+beta)^2 * (alpha_d+beta +1)  ),
         std = sqrt(var),
         succ_proba = alpha) %>%
  select(treatment, mean, std, succ_proba)

colnames(Bayes_table)=c("Treatment", "Mean", "Standard dev", "Probability optimal")

Bayes_table
#filename = paste("OUTPUT/COVID/Bayes_tables/", Sys.Date(), "_bayes_table.csv", sep = "")
#write_csv(Bayes_table, filename)
write.csv(Bayes_table,"Randomization_Inference/bayestable.csv")

# Create a file for Stata to check - if it exists, that means the R code successfully completed
cat("Done!",file="Randomization_Inference/RCodeFinished.txt",sep="\n")